TensorRT-LLM

Adam Zheng, Product Manager

目录

• 简介
• 开创性性能，现由 PyTorch 驱动
• 从可视化厨房到共享开发空间
• Roadmap
• 总结

简介

TensorRT LLM 旨在帮助用户在 NVIDIA AI 平台上获得大型语言模型（LLM）推理部署的最佳性能。

开创性性能，现由 PyTorch 驱动

TensorRT LLM 架构重构

新架构旨在通过与现有 LLM API 相同的接口，更轻松地扩展系统以支持新模型和新功能。

当前架构的挑战:
* 新模型需要手动重新实现: 需要使用 TensorRT-LLM Layer API 重写架构，并编写自定义的权重加载脚本。
* 添加新功能需要深入的系统知识: 必须理解整个紧密耦合的系统才能进行微小修改，且大多数更改必须在 C++ 中实现。
* 引擎构建限制开发速度: 在修改构建参数或升级运行时版本时，需要等待漫长的引擎构建时间。

重构后的新架构优势:
* 复用 PyTorch 实现以支持新模型: 直接复制 PyTorch 模型实现，并将关键组件替换为预构建的、包含优化内核的 torch 自定义操作（custom ops）包装器，以进一步提升性能。
* 在模块化的 Python 中构建新功能: 清晰定义的 Python 接口，用于核心运行时构建块，以实现最大的可定制性（hackability）。
* 无需构建引擎即可快速迭代: 前向传播过程在 PyTorch 中执行，无需（重新）构建引擎。

使用 TensorRT 图编译器的传统工作流程

在传统的 TensorRT 工作流程中，不能直接加载 HuggingFace 格式的 checkpoint 并运行。需要通过手动运行脚本来执行两个额外步骤：
1. 转换 Checkpoint: 准备一个与 TensorRT-LLM 中模型定义兼容的 checkpoint。
2. 构建引擎: 准备一个可运行的 TensorRT 引擎的二进制文件。

下图展示了从 HuggingFace Checkpoint 到最终运行 TRT 引擎的数据流和使用流程。

带有 PyTorch 的 LLM API

新版 LLM API 提供了基于 PyTorch 的建模支持，带来了显著的易用性提升：
* 无需转换 checkpoint: 模型定义与 vLLM/SGLang 或 HF transformers 类似，与 HF checkpoints 兼容。
* 无需构建引擎: 模型执行也基于 PyTorch。
* 统一的入口点: PyTorch 工作流的入口点也是 LLM API。
* 用户体验更好: 对用户更友好，易于修改和调试。

建模代码极大简化（以 Llama 为例）

新架构显著减少了代码量，并简化了工作流程。

旧版 TensorRT-LLM:
* tensorrt_llm/models/llama/model.py: ~600 LoC
* LLAMADecoderLayer: ~200 LoC
* LLAMAModel: ~100 LoC
* LLAMAForCausalLM: ~200 LoC
* 需要将 HF key 转换为 TRT-LLM key，量化（模块替换）等操作使代码更长。

• tensorrt_llm/models/llama/convert.py: ~2500 LoC

  • 在统一转换器之前存在旧的权重分片代码。
  • 处理具有特殊格式（如 GPTQ）的 checkpoint。

• examples/llama/convert_checkpoint.py: ~600 LoC

  • 运行 checkpoint 转换的脚本。
  • 使用 http://convert.py 中的功能。
  • 大量的命令行参数。

新版基于 PyTorch 的 TensorRT-LLM:

• tensorrt_llm/_torch/models/modeling_llama.py: ~200 LoC

  • LlamaAttention: ~30 LoC
  • LlamaDecoderLayer: ~60 LoC
  • LlamaModel: ~50 LoC
  • LlamaForCausalLM: ~15 LoC

• 无需额外的 checkpoint 转换代码:

  • 模型定义与 HF safetensors 权重字典中的键匹配，无需额外转换。
  • 权重分片在 Linear 层中处理。
  • 训练后量化委托给 ModelOpt，TRT-LLM 只需加载量化后的 checkpoint。
  • 将模型定义与加载过程隔离。

注: 这不是一个精确的苹果对苹果的比较。例如，convert.py 中的 2500 LoC 包含许多不由 TRT 工作流使用的遗留代码。但这只是为了展示 PyTorch 工作流的代码库要干净得多。

调试体验

PyTorch 的 Eager 模式执行使调试变得更加容易。

• 在 TRT 模块中:

  • 必须将张量注册为输出。
  • 需要使用 trtllm-build --enable_debug_output 等标志来构建 TRT 引擎。
  • 为 ModelRunner 打开调试模式，并在 ModelRunner 中捕获输出张量。

• 在 PyTorch 模块中:

  • 张量可以直接打印。
  • 可以在 IDE 中使用断点。

PyTorch 路径与 TRT 路径的性能对比 - 吞吐量

下图显示，基于 PyTorch 的新路径在吞吐量性能上与传统的 TRT 路径相当。

B200 上的最佳性能 - DS R1

以下是在 8xB200 硬件上，针对 DS-R1 模型的最大吞吐量和最小延迟性能数据。

• 最大吞吐量 (Max-Throughput):

  • 输入序列长度 (ISL): 1k tokens
  • 输出序列长度 (OSL): 2k tokens
  • Batch: 3072
  • num_requests: 49152
  • 每 GPU 的 TPS: 4451

• 最小延迟 (Min-Latency):

  • 输入序列长度 (ISL): 1k tokens
  • 输出序列长度 (OSL): 2k tokens
  • MTP: Enabled
  • 每用户的 TPS: 368

TensorRT LLM 入门指南

只需几行代码即可开始使用 TensorRT LLM。

from tensorrt_llm._torch import LLM

llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")

outputs = llm.generate("Open the pod bay doors, HAL")

相关链接:

• 文档: https://nvidia.github.io/TensorRT-LLM/torch.html
• 示例: https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/pytorch
• 源码: https://github.com/NVIDIA/TensorRT-LLM/tree/main/tensorrt_llm/_torch

有用的入门脚本

• examples/pytorch/quickstart.py

  • 使用一些提示运行一个 Hugging Face LLM 模型。

• examples/pytorch/quickstart_advanced.py

  • 使用不同的运行时特性和/或 tp/pp/ep 配置运行一个 Hugging Face LLM 模型。

• examples/pytorch/quickstart_multimodal.py

  • 运行一个多模态模型。

• examples/mmlu_llmapi.py

  • 使用一个 LLM 运行 mmlu 测试。

技术细节视频 - 基于 PyTorch 的架构

更多关于新的基于 PyTorch 的 TensorRT-LLM 架构的技术细节，请观看以下视频。

从可视化厨房到共享开发空间

TensorRT-LLM 与开源

开源一直是 TensorRT-LLM 的一个组成部分。

• TensorRT-LLM 自 2023 年 10 月以来一直在 GitHub 上发布源代码。
• 然而，这存在一些限制:
  • 由于 C++ 后端的限制，一些组件必须保持闭源。
  • 大部分开发在内部一个分支上进行，然后再同步到 GitHub。
  • 这使得合并来自公共贡献者的 PR 具有挑战性。
  • TensorRT-LLM 新功能和错误修复的发布延迟更长。
  • 提交历史必须被压缩 (squashed)。
  • TensorRT-LLM 的 GitHub 仓库中没有 CI/CD。

TensorRT-LLM 的新开发模型

• 引入 TensorRT-LLM 的全新 GitHub-first 开发模型！

  • 由 TensorRT-LLM 的新 PyTorch 后端启用。
  • 大多数新功能和错误修复现在都在 https://github.com/NVIDIA/TensorRT-LLM 中进行。
  • Git 历史记录现已可见。
  • 拉取请求（PRs）直接在 GitHub 中进行检查，由 NVIDIA 工程师协助。
  • 在多种 GPU 架构上提供完整的 CI/CD 支持。

• TensorRT-LLM 的更多部分已经开源

  • 例如：批处理管理器（Batch Manager）、KV 缓存管理器（KV Cache Manager）、执行器（Executor）、MHA、XQA 等。
  • 我们正在努力开源更多内容！

呼吁贡献 (Call for Contribution)

• 请遵循使用 TensorRT-LLM 的说明，并在 GitHub issues 上留下您的评论。

  • https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/torch.md
  • https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md

• 呼吁贡献（您将获得 TensorRT-LLM 团队的支持，将您的贡献引入 GitHub 仓库）。

在 GitHub 上为 TensorRT-LLM 做贡献

重要文档

• http://CODE_OF_CONDUCT.md

  • 开发者行为准则

• http://CODING_GUIDELINES.md

  • C++ 和 Python 编码风格

• http://CONTRIBUTING.md

  • 开发者工作流程文档
  • 如何提交 PR

• TensorRT-LLM Documentation

  • 开发者文档
  • 安装、API、架构、关键特性、性能

• Building from source Code on Linux

  • 构建说明

• How to run TRT-LLM tests

  • 运行和添加测试到 CI 的说明

演练：在 GitHub 上为 TensorRT-LLM 做贡献

步骤 1: 创建一个 fork。访问 https://github.com/NVIDIA/TensorRT-LLM 并点击 "Fork"。
步骤 2: git clone --recursive [email protected]:user/TensorRT-LLM.git
步骤 3: 安装 git-lfs 并运行 git lfs install; git lfs pull
步骤 4: 创建一个开发分支: git checkout -b feature_xxx
步骤 5: 构建和安装 TensorRT-LLM 并进行您的修改。
步骤 6: 添加单元测试并测试您的修改。
步骤 7: 安装 pre-commit 并在您的修改上运行它: git add .; pre-commit
步骤 8: 推送您的修改并创建一个拉取请求（Pull Request）。
步骤 9: 一位 NVIDIA 工程师将被指派审查 PR 并启动 CI/CD。
步骤 10: 在 PR 被批准并通过 CI 后，一位 NVIDIA 员工会将其合并到主分支。

技术细节视频 - GitHub-First

此部分内容通过视频形式提供了更多技术细节。视频标题为“TensorRT-LLM Goes GitHub-First”，由 Mayank Kaushik 和 Hao Lu 主讲，日期为 2025 年 4 月 30 日。

更多直播即将到来

下表列出了即将举行的在线直播活动，涵盖了从 PyTorch 架构、社区贡献到具体模型（如 DeepSeek R1）性能优化等多个主题。

NVIDIA TRT-LLM 开发技术交流群

您可以在 TensorRT-LLM 的文档页面找到最新的技术博客和新闻。

• 技术博客 (Tech Blogs)

  • [05/23] DeepSeek R1 MTP 实现和优化
  • [05/16] 推动延迟边界：在 NVIDIA B200 GPUs 上优化 DeepSeek-R1 性能

• 最新新闻 (Latest News)

  • [05/22] Blackwell 使用 Meta's Llama 4 Maverick 打破 1,000 TPS/用户屏障
  • [04/10] TensorRT-LLM DeepSeek R1 性能基准测试最佳实践现已发布
  • [04/05] TensorRT-LLM 可以在 B200 GPUs 上以超过 40,000 tokens/秒的速度运行 Llama 4

• 请微信扫码加入“TRT-LLM”专属交流群。
• 如无法扫码入群，请添加小助手微信并在申请中备注“TRT-LLM”入群。

Roadmap

TensorRT LLM 路线图

下图展示了 TensorRT-LLM 在 2025 年 5 月至 7 月的路线图，涵盖了性能、用户与开发者体验、工作流与功能等方面的规划。

2025 年 5 月 (0.19)
- 性能: Data Parallel MLA, Eagle-MTP
- 用户与开发者体验: Open-sourced C++ runtime, [prototype] AutoDeploy for automatic model optimization
- 工作流与功能: DeepSeek v3, Customizable inference engine to Triton + TRTLLM-PyTorch example

2025 年 6 月 (0.20)
- 性能: Eagle 3 and other speculative decoding enhancements
- 用户与开发者体验: Reference disaggregated orchestrator built on top of trtllm-serve, Multi-modal embedding field with LLM API
- 工作流与功能: New models including Llama4, Qwen3, Gemma3, Dynasor-CoT and stream generation scaffolding for reasoning

2025 年 7 月+ (0.21+)
- 性能: Wide Expert Parallelism, Long context optimizations, NIXL integration (faster KV cache transfer with Dynamo)
- 用户与开发者体验: Fast load times via GDS for dynamic autoscaling, Stable LLM API with backward compatibility within major version, Full feature support + testing for core models GB200 NVL72 official support, Fast KV Cache offload to / retrieval from storage
- 工作流与功能: EPD disaggregation for multi-modal models, Accelerate reasoning trace generation in RL frameworks

注：路线图可能会变更，恕不另行通知。

最新路线图和开发进展

最新的路线图和开发进展通过 GitHub 的 Issues 进行公开跟踪，方便社区了解各项功能的当前状态。

总结

本次分享的核心要点包括：

• 最佳性能 (Best Performance)
• PyTorch 架构 (PyTorch Architecture)
• GitHub-First 开发模式
• 微信交流群 (WeChat Group)